#include "molskin.h"
#include "mcconextractor.h"

using namespace PBVP;

const unsigned char MolSkin::EXTERIOR_VOLUME =   0;
const unsigned char MolSkin::SAS_SURFACE     =  50;
const unsigned char MolSkin::SAS_VOLUME      = 100;
const unsigned char MolSkin::SES_SURFACE     = 150;
const unsigned char MolSkin::SES_VOLUME      = 175;
const unsigned char MolSkin::VDW_SURFACE     = 200;
const unsigned char MolSkin::VDW_VOLUME      = 250;

MolSkin::MolSkin(QWidget *parent)
{
	m_ProbeRadius = 1.4;
	m_Depth = 3;
	m_Dim[0] = m_Dim[1] = m_Dim[2] = 64;

	progressBar = NULL;
}

Reg3DVolumeData* MolSkin::constructSkinVolume(const QVector<Atom> &atoms)
{
	assert(atoms.size() >= 1);
	setProgress(0);
	// compute minimum and max extents of the atoms
	computeMinMaxExt(atoms);
	Reg3DVolumeData *pReg3Data = new Reg3DVolumeData(m_Dim);
	pReg3Data->setMinMaxExt(m_MinExt, m_MaxExt);
	pReg3Data->setNumOfVariables(1);
	float *data = new float[m_Dim[0]*m_Dim[1]*m_Dim[2]];
	memset(data, 0, sizeof(float)*m_Dim[0]*m_Dim[1]*m_Dim[2]);
	setProgress(5);

	m_NumCells = (m_Dim[0]-1)*(m_Dim[1]-1)*(m_Dim[2]-1);
	m_NumVerts = m_Dim[0]*m_Dim[1]*m_Dim[2];
	//m_NumSmallCells = pow(pow(2.0, m_Depth-1), 3);
	m_NumSmallCells = (1 << (m_Depth-1))*(1 << (m_Depth-1))*(1 << (m_Depth-1));
	//m_NumSmallVerts = pow(pow(2.0, m_Depth-1)+1, 3);
	m_NumSmallVerts = ((1<<(m_Depth-1))+1)*((1<<(m_Depth-1))+1)*((1<<(m_Depth-1))+1);

	m_SasInteriorVertices = new bool[m_NumVerts];
	m_SasBoundaryVertices = new bool[m_NumVerts];
	m_VDWInteriorVertices = new bool[m_NumVerts];
	m_VDWBoundaryVertices = new bool[m_NumVerts];

	m_SasBoundaryCells = new bool[m_NumCells];

	m_IndexOfBoundaryCell = new int[m_NumCells];

	m_Regions = new unsigned char[m_NumVerts];

	for(int i = 0; i < m_NumVerts; i++) {
		m_SasInteriorVertices[i] = false;
		m_SasBoundaryVertices[i] = false;
		m_VDWInteriorVertices[i] = false;
		m_VDWBoundaryVertices[i] = false;
		m_Regions[i] = EXTERIOR_VOLUME;
		data[i] = 0;
	}

	getSolventInterior(atoms);
	getBoundaryCells();
	setProgress(40);

	m_NumBoundaryCells = 0;
	for(int i = 0; i < m_NumCells; i++ )
	{
		if( m_SasBoundaryCells[i] )
		{
			m_IndexOfBoundaryCell[i] = m_NumBoundaryCells;
			m_NumBoundaryCells++;
		}
		else
			m_IndexOfBoundaryCell[i] = -1;
	}

	m_SphereIndices = new std::vector<int>[m_NumBoundaryCells];
	m_SubdividedNodeSasInteriorVertex = new bool[m_NumBoundaryCells*m_NumSmallVerts];
	m_SubdividedNodeSasBoundaryCell = new bool[m_NumBoundaryCells*m_NumSmallCells];

	getBoundarySphereIndices(atoms);
	setProgress(50);
	getSubdividedNodesSasInteriorVertices(atoms);
	setProgress(60);
	getSubdividedNodesSasBoundaryCells();
	setProgress(70);

	//getSesSurfaceAndVolume();

	int imin=0, jmin=0, kmin=0, imax=0, jmax=0, kmax=0;
	int iProbeRadius = ceil(m_ProbeRadius/m_Span[0]);
	int jProbeRadius = ceil(m_ProbeRadius/m_Span[1]);
	int kProbeRadius = ceil(m_ProbeRadius/m_Span[2]);

	for(int v=0; v < m_NumVerts; v++ )
	{
		if( m_Regions[v] == EXTERIOR_VOLUME ) data[v] = 0;
		
		if( m_Regions[v] == VDW_VOLUME || m_Regions[v] == EXTERIOR_VOLUME ) 
		{
			continue;
		}

		double minDistToSas = 1.0e8;
		int x, y, z;
		x = v % m_Dim[0];
		y = (v / m_Dim[0]) % m_Dim[1];
		z = v /(m_Dim[0]*m_Dim[1]);

		imin = (x - iProbeRadius - 1 < 0)? 0 : x - iProbeRadius - 1;
		jmin = (y - jProbeRadius - 1 < 0)? 0 : y - jProbeRadius - 1; 
		kmin = (z - kProbeRadius - 1 < 0)? 0 : z - kProbeRadius - 1; 

		imax = (x + iProbeRadius + 1 > m_Dim[0]-1 )? m_Dim[0]-1 : x + iProbeRadius + 1;
		jmax = (y + jProbeRadius + 1 > m_Dim[1]-1 )? m_Dim[1]-1 : y + jProbeRadius + 1;
		kmax = (z + kProbeRadius + 1 > m_Dim[2]-1 )? m_Dim[2]-1 : z + kProbeRadius + 1;
		
		for(int i = imin; i < imax; i++ )
		{
			for(int j = jmin; j < jmax; j++ )
			{
				for(int k = kmin; k < kmax; k++ )
				{
					int cell = i + j*(m_Dim[0]-1) + k*(m_Dim[0]-1)*(m_Dim[1]-1);
					if( !m_SasBoundaryCells[cell] ) continue;

					double curMinDistToSas = getDistToSas( x, y, z, i, j, k );
					if( minDistToSas > curMinDistToSas )
						minDistToSas = curMinDistToSas;
				}
			}
		}
		if( minDistToSas > m_ProbeRadius )
		{
			m_Regions[v] = SES_VOLUME;
			data[v] = minDistToSas;
			if( minDistToSas > 4*m_ProbeRadius )  data[v] = 4*m_ProbeRadius;
		}
		else
		{
			double scaleFactor = (VDW_SURFACE - SAS_SURFACE ) / (4.0*m_ProbeRadius);
			double dist = minDistToSas * scaleFactor;
			if( dist > VDW_SURFACE )
			{
				dist = VDW_SURFACE;
			}

			m_Regions[v] = SAS_SURFACE + dist;
			data[v] = minDistToSas; // isosurface = 87.5
		}
	}
	pReg3Data->setDataArray(data);
	setProgress(90);

	delete[] m_SasInteriorVertices;
	delete[] m_SasBoundaryVertices;
	delete[] m_VDWInteriorVertices;
	delete[] m_VDWBoundaryVertices;
	delete[] m_SasBoundaryCells;
	delete[] m_IndexOfBoundaryCell;
	delete[] m_Regions;
	delete[] m_SphereIndices;
	delete[] m_SubdividedNodeSasInteriorVertex;
	delete[] m_SubdividedNodeSasBoundaryCell;

	return pReg3Data;
}

void MolSkin::computeMinMaxExt(const QVector<Atom> &atoms)
{
	m_MinExt[0] = m_MaxExt[0] = atoms[0].position[0];
	m_MinExt[1] = m_MaxExt[1] = atoms[0].position[1];
	m_MinExt[2] = m_MaxExt[2] = atoms[0].position[2];
	for (int i = 1; i < atoms.size(); i++)
	{
		if (m_MinExt[0] > atoms[i].position[0]) m_MinExt[0] = atoms[i].position[0];
		if (m_MinExt[1] > atoms[i].position[1]) m_MinExt[1] = atoms[i].position[1];
		if (m_MinExt[2] > atoms[i].position[2]) m_MinExt[2] = atoms[i].position[2];
		if (m_MaxExt[0] < atoms[i].position[0]) m_MaxExt[0] = atoms[i].position[0];
		if (m_MaxExt[1] < atoms[i].position[1]) m_MaxExt[1] = atoms[i].position[1];
		if (m_MaxExt[2] < atoms[i].position[2]) m_MaxExt[2] = atoms[i].position[2];
	}
	double offset = 4*m_ProbeRadius;
	m_MinExt[0] -= offset; 
	m_MinExt[1] -= offset; 
	m_MinExt[2] -= offset;
	m_MaxExt[0] += offset; 
	m_MaxExt[1] += offset;
	m_MaxExt[2] += offset;
	m_Orig[0] = m_MinExt[0];
	m_Orig[1] = m_MinExt[1];
	m_Orig[2] = m_MinExt[2];
	m_Span[0] = (m_MaxExt[0]-m_MinExt[0])/(m_Dim[0]-1);
	m_Span[1] = (m_MaxExt[1]-m_MinExt[1])/(m_Dim[1]-1);
	m_Span[2] = (m_MaxExt[2]-m_MinExt[2])/(m_Dim[2]-1);
}

void MolSkin::getSolventInterior(const QVector<Atom> &atoms)
{
	for(int k = 0; k < atoms.size(); k++) {
		Atom atom = atoms[k];
		double radius = atom.radius;
		Vector3 center = atom.position;
		double maxRadius = radius + m_ProbeRadius + 1;
		int minIndex[3], maxIndex[3];
		for(int i = 0; i < 3; i++) {
			minIndex[i] = (int)((center[i]-m_MinExt[i]-maxRadius)/m_Span[i] - 0.5);
			minIndex[i] = (minIndex[i] < 0)? 0 : minIndex[i];
			maxIndex[i] = (int)((center[i]-m_MinExt[i]+maxRadius)/m_Span[i] + 0.5);
			maxIndex[i] = (maxIndex[i] > (m_Dim[i]-1))? (m_Dim[i]-1) : maxIndex[i];
		}
		for(int iz = minIndex[2]; iz <= maxIndex[2]; iz++)
			for(int iy = minIndex[1]; iy <= maxIndex[1]; iy++)
				for (int ix = minIndex[0]; ix <= maxIndex[0]; ix++) 
				{
					int nv = ix + iy*m_Dim[0] + iz*m_Dim[0]*m_Dim[1];
					Vector3 pos(m_Orig[0]+ix*m_Span[0], m_Orig[1]+iy*m_Span[1], m_Orig[2]+iz*m_Span[2]);
					double r_2 = (center[0]-pos[0])*(center[0]-pos[0]) +
								 (center[1]-pos[1])*(center[1]-pos[1]) +
								 (center[2]-pos[2])*(center[2]-pos[2]);
					if(r_2 < (radius+m_ProbeRadius)*(radius+m_ProbeRadius)) 
					{
						m_SasInteriorVertices[nv] = true;
						if(m_Regions[nv] != VDW_VOLUME) {
							m_Regions[nv] = SAS_VOLUME;
						}
					}
					if (r_2 < (radius*radius))
					{
						m_VDWInteriorVertices[nv] = true;
						m_Regions[nv] = VDW_VOLUME;
					} 
				}
	}
}

/*!
Once vertices are classified as interior or not, we find and mark the boundary cells and mark all vertices
belonging to a boundary cell as a boundary vertex

Also first classify m_Regions as VDW boundary vertices and then as SAS boundary vertices. This priority was
chosen arbitrarily
*/
void MolSkin::getBoundaryCells()
{
	for (unsigned int k = 0; k < m_Dim[2]-1; k++) 
		for(unsigned int j = 0; j < m_Dim[1]-1; j++)
			for (unsigned int i = 0; i < m_Dim[0]-1; i++)
	{
		int nc = i + j*(m_Dim[0]-1) + k*(m_Dim[0]-1)*(m_Dim[1]-1);
		int vertId[8];
		getCellVerts(i, j, k, vertId, m_Dim);
		if (isBoundaryCell(i, j, k, m_VDWInteriorVertices))
		{
			for(int v = 0; v < 8; v++) {
				m_Regions[vertId[v]] = VDW_SURFACE;
				m_VDWBoundaryVertices[vertId[v]] = true;
			}
		}
		if (isBoundaryCell(i, j, k, m_SasInteriorVertices)) 
		{
			for (int v = 0; v < 8; v++)
			{
				if (m_Regions[vertId[v]] == SAS_VOLUME)
				{
					m_Regions[vertId[v]] = SAS_SURFACE;
					m_SasBoundaryVertices[vertId[v]] = true;
				}
			}
		}
		m_SasBoundaryCells[nc] = true;
	}
	
}

bool MolSkin::isBoundaryCell(int i, int j, int k, bool *interior)
{
	bool vertFlags[8];
	int v = i + j*m_Dim[0] + k*m_Dim[0]*m_Dim[1];
	vertFlags[0] = interior[v];
	vertFlags[1] = interior[v + 1];
	vertFlags[2] = interior[v + m_Dim[0]*m_Dim[1] + 1];
	vertFlags[3] = interior[v + m_Dim[0]*m_Dim[1]];
	vertFlags[4] = interior[v + m_Dim[0]];
	vertFlags[5] = interior[v + m_Dim[0] + 1];
	vertFlags[6] = interior[v + m_Dim[0]*m_Dim[1] + m_Dim[0] + 1];
	vertFlags[7] = interior[v + m_Dim[0]*m_Dim[1] + m_Dim[0]];

	bool inFlag = false, outFlag = false;
	for(int a = 0; a < 8; a++) {
		if (vertFlags[a]) inFlag = true;
		else outFlag = true;
	}
	return inFlag & outFlag;
}

bool MolSkin::isBoundaryCell(bool* interior)
{
	bool interiorFlag,exteriorFlag;
	interiorFlag=exteriorFlag=false;

	for (int i=0;i<8;i++) {
		if (interior[i]==true) interiorFlag=true;
		else exteriorFlag=true;
	}

	return (interiorFlag&exteriorFlag);
}

void MolSkin::getCellVerts(int i, int j, int k, int vertId[8], unsigned int dim[3])
{
	int v = i + j*dim[0] + k*dim[0]*dim[1];
	vertId[0] = v;
	vertId[1] = v + 1;
	vertId[2] = v + dim[0]*dim[1] + 1;
	vertId[3] = v + dim[0]*dim[1];
	vertId[4] = v + dim[0];
	vertId[5] = v + dim[0] + 1;
	vertId[6] = v + dim[0]*dim[1] + dim[0] + 1;
	vertId[7] = v + dim[0]*dim[1] + dim[0];
}

/*!
Find the distance of the grid point vertex to the SAS. Search in the small cells of the given cell
Take the distance to be the distance to the center of the small cell.
*/
double MolSkin::getDistToSas( int vx, int vy, int vz, int cx, int cy, int cz )
{
	double curDist = 1.0e8;

	int cell = cx + cy*(m_Dim[0]-1) + cz*(m_Dim[0]-1)*(m_Dim[1]-1);
	int cellIndex = m_IndexOfBoundaryCell[cell];

	unsigned int dim[3];
	//dim[0] = dim[1] = dim[2] = pow(2, m_Depth-1)+1;
	dim[0] = dim[1] = dim[2] = (1 << (m_Depth-1)) + 1;

	for(int i=0; i < m_NumSmallCells; i++ )
	{
		int smallCellIndex = cellIndex*m_NumSmallCells + i;
		if( !m_SubdividedNodeSasBoundaryCell[smallCellIndex] ) continue;

		int cellx, celly, cellz;
		cellx = i % (dim[0]-1);
		celly = (i / (dim[0]-1)) % (dim[1]-1);
		cellz = i / ((dim[0]-1)*(dim[1]-1));
		// TrilinearGrid::cell2xyz(i, cellx, celly, cellz, dim );

		if( (cellx+0.5)/((double)dim[0]) > 1.01 )
		{
			printf("Error\n");
		}
		if( (celly+0.5)/((double)dim[1]) > 1.01 )
		{
			printf("Error\n");
		}
		if( (cellz+0.5)/((double)dim[2]) > 1.01 )
		{
			printf("Error\n");
		}
		double d = getDist(	vx, vy, vz, 
			cx+(cellx+0.5)/((double)dim[0]), 
			cy+(celly+0.5)/((double)dim[1]), 
			cz+(cellz+0.5)/((double)dim[2]) );

		if( curDist > d ) curDist = d;
	}
	return curDist;
}

/**
 * Given two grid points, compute the distance between the two of them in real space
 */
double MolSkin::getDist( double x1, double y1, double z1, double x2, double y2, double z2 )
{
	double x = fabs(x1-x2)*m_Span[0];
	double y = fabs(y1-y2)*m_Span[1];
	double z = fabs(z1-z2)*m_Span[2];

	return sqrt(x*x + y*y + z*z);
}

/*!
For each sphere, find all the cells it intersects.
If it intersects a boundary cell, add its index to the m_SphereIndices array at the right location
*/
void MolSkin::getBoundarySphereIndices(const QVector<Atom>& atoms)
{
	int numAdded = 0;
	int numBCells = 0;

	for(int s=0; s < atoms.size(); s++ )
	{
		Atom atom = atoms[s];
		double radius = atom.radius;
		Vector3 center = atom.position;

		int bbmin[3], bbmax[3];

		for( int i=0; i<3; i++ )
		{
			bbmin[i] = floor((center[i]-(radius+m_ProbeRadius)-m_Orig[i]) / m_Span[i]);
			bbmin[i]--;
			if( bbmin[i] < 0 ) bbmin[i] = 0;

			bbmax[i] = floor((center[i]+(radius+m_ProbeRadius)-m_Orig[i]) / m_Span[i]);// we take floor because we only need the lower vertex for the cell
			bbmax[i]++;
			if( bbmax[i] > m_Dim[i]-2 ) bbmax[i] = m_Dim[i]-2;
		}


		for(int i = bbmin[0]; i <= bbmax[0]; i++ )
		{
			for(int j = bbmin[1]; j <= bbmax[1]; j++ )
			{
				for(int k = bbmin[2]; k <= bbmax[2]; k++ )
				{
					int c = i + j*(m_Dim[0]-1) + k*(m_Dim[0]-1)*(m_Dim[1]-1);
					if( !m_SasBoundaryCells[c] ) continue;

					numBCells++;

					//if( intersects( center, radius+m_ProbeRadius, c ) )
					if( sphereCellIntersect(center, radius+m_ProbeRadius, i, j, k ) )
					{
						m_SphereIndices[m_IndexOfBoundaryCell[c]].push_back( s );
						numAdded++;
					}
				}
			}
		}
	}
	printf("Number added = %d\n", numAdded );
}

/** 
 * Check if the cell (ix, iy, iz) is inside or intersects the sphere
 * at center with the given radius
 */
bool MolSkin::sphereCellIntersect( const Vector3& center, double radius, int ix, int iy, int iz )
{
	double Bmin[3], Bmax[3];
	Bmin[0] = m_Orig[0] + m_Span[0] * ix;
	Bmin[1] = m_Orig[1] + m_Span[1] * iy;
	Bmin[2] = m_Orig[2] + m_Span[2] * iz;

	Bmax[0] = Bmin[0] + m_Span[0];
	Bmax[1] = Bmin[1] + m_Span[1];
	Bmax[2] = Bmin[2] + m_Span[2];

	double dmin = 0;
	for(int i = 0; i < 3; i++ ) 
	{	// this code is problematic (xiaoyu)
		/* if( center[i] < Bmin[i] ) 
			dmin += (center[i] - Bmin[i] )*(center[i] - Bmin[i] ); 
		else if( center[i] > Bmax[i] ) 
			dmin += ( center[i] - Bmax[i] )*( center[i] - Bmax[i] );     
		*/
		double d1 = (center[i] - Bmin[i] )*(center[i] - Bmin[i] );
		double d2 = ( center[i] - Bmax[i] )*( center[i] - Bmax[i]);
		if(d1 < d2) dmin += d1;
		else dmin += d2;
	}

	if( dmin <= radius*radius ) return true;

	return false;
}

/**
For each boundary cell, we have the sphere that pass through it. 
For each sphere in a cell, we find the subvertices of the cell that are interior to it.
*/
void MolSkin::getSubdividedNodesSasInteriorVertices(const QVector<Atom>& atoms)
{
	//int n = pow( 2.0, (m_Depth-1) )+1; // number of vertices in each side
	int n = (1 << (m_Depth-1)) + 1;

	int nBCV = m_NumBoundaryCells*m_NumSmallVerts;
	for(int i=0; i < nBCV; i++ ){
		m_SubdividedNodeSasInteriorVertex[i] = false;
	}

	unsigned int small_dim[3];
	small_dim[0] = small_dim[1] = small_dim[2] = n;
	for(int c = 0; c < m_NumCells; c++ )
	{
		if( !m_SasBoundaryCells[c] ) continue;
		int cellIndex = m_IndexOfBoundaryCell[c];

		int celli, cellj, cellk;
		celli = c % (m_Dim[0]-1);
		cellj = (c / (m_Dim[0]-1))%(m_Dim[1]-1);
		cellk = c / ((m_Dim[0]-1)*(m_Dim[1]-1));

		for(unsigned int s=0; s < m_SphereIndices[cellIndex].size(); s++ )
		{
			Atom atom = atoms[m_SphereIndices[cellIndex].at(s)];
			double radius = atom.radius; 
			radius += m_ProbeRadius;

			double sx = atom.position[0];
			double sy = atom.position[1];
			double sz = atom.position[2];

			for(unsigned int sv = 0; sv < m_NumSmallVerts; sv++ )
			{
				int i, j, k;
				i = sv % small_dim[0];
				j = (sv / small_dim[0])%small_dim[1];
				k = sv / (small_dim[0]*small_dim[1]);

				double vx = m_Orig[0] + (celli+i/((double)n)) * m_Span[0];
				double vy = m_Orig[1] + (cellj+j/((double)n)) * m_Span[1];
				double vz = m_Orig[2] + (cellk+k/((double)n)) * m_Span[2];

				double dist = (sx-vx)*(sx-vx) + (sy-vy)*(sy-vy) + (sz-vz)*(sz-vz);
				if( dist < radius*radius )
					m_SubdividedNodeSasInteriorVertex[cellIndex*m_NumSmallVerts + sv] = true;
			}
		}
	}
}

/**
 * We already know which small vertices are in and out. 
 *We now classify the cells as boundary or not.
 */
void MolSkin::getSubdividedNodesSasBoundaryCells()
{
	//int n = pow( 2.0, (m_Depth-1) )+1; // number of small verts in each side
	int n = (1 << (m_Depth-1)) + 1;
	unsigned int small_dim[3];
	small_dim[0] = small_dim[1] = small_dim[2] = n;

	int nBC = m_NumBoundaryCells*m_NumSmallCells;
	for(int i=0; i < nBC; i++ ) {
		m_SubdividedNodeSasBoundaryCell[i] = false;
	}
	
	for(int c = 0; c < m_NumCells; c++ )
	{
		if( !m_SasBoundaryCells[c] ) continue;
		int cellIndex = m_IndexOfBoundaryCell[c];

		int celli, cellj, cellk;
		celli = c % (m_Dim[0]-1);
		cellj = (c / (m_Dim[0]-1))%(m_Dim[1]-1);
		cellk = c / ((m_Dim[0]-1)*(m_Dim[1]-1));

		//for(int sc = 0; sc < m_NumSmallCells; sc++ )
		for (int i = 0; i < small_dim[0]-1; i++)
			for(int j = 0; j < small_dim[1]-1; j++)
				for(int k = 0; k < small_dim[2]-1; k++)
		{
			int sc = i + j*(small_dim[0]-1) + k*(small_dim[0]-1)*(small_dim[1]-1);
			int cellVertexArray[8];
			getCellVerts(i, j, k, cellVertexArray, small_dim);
			// TrilinearGrid::getCellVertices(sc,cellVertexArray,small_dim);

			bool sasCellInterior[8];
			for (int l =0; l < 8; l++)
				sasCellInterior[l]=m_SubdividedNodeSasInteriorVertex[cellIndex*m_NumSmallVerts + cellVertexArray[l]];

			if (isBoundaryCell(sasCellInterior)) 
				m_SubdividedNodeSasBoundaryCell[cellIndex*m_NumSmallCells + sc] = true;
		}
	}	
}

GeometryData* MolSkin::constructMolSurface(const QVector<Atom> &atoms)
{
	Reg3DVolumeData* pReg3Data = constructSkinVolume(atoms);
	MCContourExtractor extractor;
	GeometryData *pContour = extractor.extractContour(pReg3Data, m_ProbeRadius);
	setProgress(100);
	delete pReg3Data;
	return pContour;
}

void MolSkin::setProgress(int prog)
{
	if(progressBar) {
		progressBar->setValue(prog);
	}
}